<?php

namespace collections\map;

use collections\compare\ComparableObject;
use collections\compare\Comparator;
use collections\map\TreeEntry as E;

class TreeMap extends AbstractMap implements \ArrayAccess
{
    /** @var Comparator */
    private $comparator;

    /** @var TreeEntry */
    private $root;

    /** @var int */
    private $size;

    /** @var int */
    private $modCount;

    public function __construct($comparator = null) {
        $this->comparator = $comparator;
        $this->size = $this->modCount = 0;
    }
    public function size() { return $this->size; }

    public function getEntry(/* mixed */$key) {
        $p = $this->root;
        while ($p != null) {
            if (!is_null($this->comparator)) {
                $cmp = $this->comparator->compare($key, $p->getKey());
            } else {
                $k = new ComparableObject($key);
                $cmp = $k->compareTo($p->getKey());
            }
            if ($cmp < 0) {
                $p = $p->left;
            } else if ($cmp > 0) {
                $p = $p->right;
            } else {
                return $p;
            }
        }
        return null;
    }
    public function getFirstEntry() {
        $p = $this->root;
        if (!is_null($p)) {
            while (!is_null($p->left)) {$p = $p->left;}
        }
        return $p;
    }
    public function getLastEntry() {
        $p = $this->root;
        if (!is_null($p)) {
            while (!is_null($p->right)) {$p = $p->right;}
        }
        return $p;
    }

    public function get($key) {
        $p = $this->getEntry($key);
        return is_null($p) ? null : $p->value;
    }

    public function put($key, $value) {
        $t = $this->root;
        if (is_null($t)) {
            $this->compare($key, $key); // type (and possibly null) check
            $this->root = new TreeEntry($key, $value, null);
            $this->size = 1;
            $this->modCount += 1;
            return null;
        }
        /** @var $cpr Comparator */
        $cpr = $this->comparator;

        if (!is_null($cpr)) {
            do {
                $parent = $t;
                $cmp = $cpr->compare($key, $t->key);
                if ($cmp < 0) {
                    $t = $t->left;
                } elseif ($cmp > 0) {
                    $t = $t->right;
                } else {
                    return $t->setValue($value);
                }
            } while(!is_null($t));
        } else {
            if (is_null($key)) {throw new \InvalidArgumentException("Invalid entry key");}
            do {
                $parent = $t;
                $k = new ComparableObject($key);
                $tk = new ComparableObject($t->key);
                $cmp = $k->compareTo($tk);
                if ($cmp < 0) {
                    $t = $t->left;
                } else if ($cmp > 0) {
                    $t= $t->right;
                } else {
                    return $t->setValue($value);
                }
            } while (!is_null($t));
        }
        $e = new TreeEntry($key, $value, $parent);
        if ($cmp < 0) {
            $parent->left = $e;
        } else {
            $parent->right = $e;
        }
        $this->fixAfterInsertion($e);
        $this->size += 1;
        $this->modCount += 1;
        return null;
    }

    public function remove(/* mixed */$key) {
        $p = $this->getEntry($key);
        if (is_null($p)) {return null;}
        $oldValue = $p->getValue();
        $this->deleteEntry($p);
        return $oldValue;
    }

    private function compare($a, $b) {
        if (empty($this->comparator)) {
            $a0 = new ComparableObject($a);
            $b0 = new ComparableObject($b);
            return $a0->compareTo($b0);
        }
        return $this->comparator->compare($a, $b);
    }

    /**
     * @param $p TreeEntry
     */
    public function rotateLeft($p) {
        if (is_null($p)) {return;}
        $r = $p->right;
        $p->right = $r->left;
        if (!is_null($r->left)) {
            $r->left->parent = $p;
        }
        $r->parent = $p->parent;
        if (is_null($p->parent)) {
            $this->root = $r;
        } elseif ($p->parent->left === $p) {
            $p->parent->left = $r;
        } else {
            $p->parent->right = $r;
        }
        $r->left = $p;
        $p->parent = $r;
    }

    /** @param $p TreeEntry */
    public function rotateRight($p) {
        if (is_null($p)) {return;}
        $l = $p->left;
        $p->left = $l->right;
        if (!is_null($l->right)) {
            $l->right->parent = $p;
        }
        $l->parent = $p->parent;
        if (is_null($p->parent)) {
            $this->root = $l;
        } elseif ($p->parent->right === $p) {
            $p->parent->right = $l;
        } else {
            $p->parent->left = $l;
        }
        $l->right = $p;
        $p->parent = $l;
    }

    private function fixAfterInsertion(E $x) {
        $x->color = E::RED;
        while (!is_null($x) && $x !== $this->root && $x->parent->color == E::RED) {
            if (E::parentOf($x) === E::leftOf(E::parentOf(E::parentOf($x)))) {
                $y = E::rightOf(E::parentOf(E::parentOf($x)));
                if (E::colorOf($y) == E::RED) {
                    E::setColor(E::parentOf($x), E::BLACK);
                    E::setColor($y, E::BLACK);
                    E::setColor(E::parentOf(E::parentOf($x)), E::RED);
                    $x = E::parentOf(E::parentOf($x));
                } else {
                    if ($x === E::rightOf(E::parentOf($x))) {
                        $x = E::parentOf($x);
                        $this->rotateLeft($x);
                    }
                    E::setColor(E::parentOf($x), E::BLACK);
                    E::setColor(E::parentOf(E::parentOf($x)), E::RED);
                    $this->rotateRight(E::parentOf(E::parentOf($x)));
                }
            } else {
                $y = E::leftOf(E::parentOf(E::parentOf($x)));
                if (E::colorOf($y) == E::RED) {
                    E::setColor(E::parentOf($x), E::BLACK);
                    E::setColor($y, E::BLACK);
                    E::setColor(E::parentOf(E::parentOf($x)), E::RED);
                    $x = E::parentOf(E::parentOf($x));
                } else {
                    if ($x === E::leftOf(E::parentOf($x))) {
                        $x = E::parentOf($x);
                        $this->rotateRight($x);
                    }
                    E::setColor(E::parentOf($x), E::BLACK);
                    E::setColor(E::parentOf(E::parentOf($x)), E::RED);
                    $this->rotateLeft(E::parentOf(E::parentOf($x)));
                }
            }
        }
        $this->root->color = E::BLACK;
    }

    /**
     * Delete node p, and then re-balance the tree.
     * @param $p TreeEntry
     */
    private function deleteEntry($p) {
        $this->modCount += 1;
        $this->size -= 1;
        // If strictly internal, copy successor's element to p and then make p
        // point to successor.
        if (!is_null($p->left) && !is_null($p->right)) {
            $s = TreeEntry::successor($p);
            $p->key = $s->key;
            $p->value = $s->value;
            $p = $s;
        } // p has 2 children
        // Start fixup at replacement node, if it exists.
        $replacement = (!is_null($p->left) ? $p->left : $p->right);
        if (!is_null($replacement)) {
            // Link replacement to parent
            $replacement->parent = $p->parent;
            if (is_null($p->parent))
                $this->root = $replacement;
            else if ($p === $p->parent->left)
                $p->parent->left  = $replacement;
            else
                $p->parent->right = $replacement;

            // Null out links so they are OK to use by fixAfterDeletion.
            $p->left = $p->right = $p->parent = null;

            // Fix replacement
            if ($p->color === TreeEntry::BLACK) {
                $this->fixAfterDeletion($replacement);
            }
        } else if (is_null($p->parent)) { // return if we are the only node.
            $this->root = null;
        } else { //  No children. Use self as phantom replacement and unlink.
            if ($p->color == TreeEntry::BLACK)
                $this->fixAfterDeletion($p);

            if (!is_null($p->parent)) {
                if ($p === $p->parent->left) {
                    $p->parent->left = null;
                } elseif ($p === $p->parent->right) {
                    $p->parent->right = null;
                }
                $p->parent = null;
            }
        }
    }

    /**
     * @param $x TreeEntry
     */
    private function fixAfterDeletion($x) {
        while ($x !== $this->root && E::colorOf($x)== E::BLACK) {
            if ($x === E::leftOf(E::parentOf($x))) {
                $sib = E::rightOf(E::parentOf($x));
                if (E::colorOf($sib) == E::RED) {
                    E::setColor($sib, E::BLACK);
                    E::setColor(E::parentOf($x), E::RED);
                    $this->rotateLeft(E::parentOf($x));
                    $sib = E::rightOf(E::parentOf($x));
                }

                if (E::colorOf(E::leftOf($sib))  == E::BLACK &&
                    E::colorOf(E::rightOf($sib)) == E::BLACK) {
                    E::setColor($sib, E::RED);
                    $x = E::parentOf($x);
                } else {
                    if (E::colorOf(E::rightOf($sib)) == E::BLACK) {
                        E::setColor(E::leftOf($sib), E::BLACK);
                        E::setColor($sib, E::RED);
                        $this->rotateRight($sib);
                        $sib = E::rightOf(E::parentOf($x));
                    }
                    E::setColor($sib, E::colorOf(E::parentOf($x)));
                    E::setColor(E::parentOf($x), E::BLACK);
                    E::setColor(E::rightOf($sib), E::BLACK);
                    $this->rotateLeft(E::parentOf($x));
                    $x = $this->root;
                }
            } else {
                $sib = E::leftOf(E::parentOf($x));
                if (E::colorOf($sib) === E::RED) {
                    E::setColor($sib, E::BLACK);
                    E::setColor(E::parentOf($x), E::RED);
                    $this->rotateRight(E::parentOf($x));
                    $sib = E::leftOf(E::parentOf($x));
                }

                if (E::colorOf(E::rightOf($sib))  == E::BLACK &&
                    E::colorOf(E::leftOf($sib)) == E::BLACK) {
                    E::setColor($sib, E::RED);
                    $x = E::parentOf($x);
                } else {
                    if (E::colorOf(E::leftOf($sib)) == E::BLACK) {
                        E::setColor(E::rightOf($sib), E::BLACK);
                        E::setColor($sib, E::RED);
                        $this->rotateLeft($sib);
                        $sib = E::leftOf(E::parentOf($x));
                    }
                    E::setColor($sib, E::colorOf(E::parentOf($x)));
                    E::setColor(E::parentOf($x), E::BLACK);
                    E::setColor(E::leftOf($sib), E::BLACK);
                    $this->rotateRight(E::parentOf($x));
                    $x = $this->root;
                }
            }
        }
        E::setColor($x, E::BLACK);
    }

    public function entrySet()
    {
        $entrySet = array();
        /** @var $e TreeEntry */
        for ($e = $this->getFirstEntry(); $e != null; $e = TreeEntry::successor($e)) {
            array_push($entrySet, $e);
        }
        return $entrySet;
    }

    public function containsKey(/* mixed */ $key) {
        $entry = $this->getEntry($key);
        return !is_null($entry);
    }

    public function containsValue($value)
    {
        /** @var $e TreeEntry */
        for ($e = $this->getFirstEntry(); $e != null; $e = TreeEntry::successor($e)) {
            if (TreeEntry::valEquals($value, $e->value)) {
                return true;
            }
        }
        return false;
    }
    public function each(callable $action) {
        for ($e = $this->getFirstEntry(); $e != null; $e = TreeEntry::successor($e)) {
            call_user_func($action, $e->getKey(), $e->getValue());
        }
    }

    public function offsetExists($offset) { return $this->containsKey($offset); }

    public function offsetGet($offset) {return $this->get($offset);}

    public function offsetSet($offset, $value) { return $this->put($offset, $value); }

    public function offsetUnset($offset) {
        $entry = $this->getEntry($offset);
        if ($entry) {
            $this->deleteEntry($entry);
        }
    }
}